package org.hhf.ai.tools;

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
public class SSEUtils {
    // timeout
    public static Long DEFAULT_TIME_OUT = 120_000L;
    // 订阅表
    private static Map<String, SseEmitter> subscribeMap = new ConcurrentHashMap<>();

    /** 添加订阅 */
    public static SseEmitter create(String sessionId) {
        if (null == sessionId || "".equals(sessionId)) {
            return null;
        }

        SseEmitter emitter = subscribeMap.get(sessionId);
        if (null == emitter) {
            emitter = new SseEmitter(DEFAULT_TIME_OUT);
            subscribeMap.put(sessionId, emitter);
        }

        /*emitter.onCompletion(() -> {
            log.info("结束连接：{}", sessionId);
            close(sessionId);
        });
        emitter.onError(throwable -> {
            log.info("连接异常：{}", sessionId);
            close(sessionId);
        });
        emitter.onTimeout(() -> {
            log.info("连接超时：{}", sessionId);
            close(sessionId);
        });*/
        return emitter;
    }

    /** 发消息 */
    public static void send(String sessionId, Object msg) {
        SseEmitter emitter = subscribeMap.get(sessionId);
        if (null != emitter) {
            try {
                emitter.send(msg);
//                emitter.send(SseEmitter.event().data(msg));
            } catch (Exception e) {
                log.error(e.getMessage(), e);
                emitter.completeWithError(e);
            }
        }
    }

    /**
     * 关闭订阅
     * @param sessionId
     */
    public static void close(String sessionId) {
        SseEmitter emitter = subscribeMap.get(sessionId);
        if (null != emitter) {
            try {
                emitter.complete();
                subscribeMap.remove(sessionId);
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
    }

    /**
     * 获取所有连接
     * @return
     */
    public static Map<String, SseEmitter> getConnections() {
        return subscribeMap;
    }
}